package rabbit.flt.plugins.spring.r2dbc.plugin;

import org.springframework.r2dbc.core.FetchSpec;
import rabbit.flt.plugins.common.R2dbcSupportPlugin;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.lang.reflect.Method;
import java.util.function.Supplier;

public class ExecuteSpecSupportPlugin extends R2dbcSupportPlugin {

    @Override
    public Object after(Object objectEnhanced, Method method, Object[] args, Object result) {
        super.after(objectEnhanced, method, args, result);
        if (isTraceOpened()) {
            Supplier<String> sqlSupplier = (Supplier<String>) args[0];
            String sql = sqlSupplier.get().trim();
            FetchSpec highVersionSpec = (FetchSpec) result;
            return new FetchSpec() {

                @Override
                public Mono<Integer> rowsUpdated() {
                    return (Mono<Integer>) wrapResult(highVersionSpec.rowsUpdated(), method, sql);
                }

                @Override
                public Mono one() {
                    return wrapResult(highVersionSpec.one(), method, sql);
                }

                @Override
                public Mono first() {
                    return wrapResult(highVersionSpec.first(), method, sql);
                }

                @Override
                public Flux all() {
                    return wrapResult(highVersionSpec.all(), method, sql);
                }
            };
        }
        return result;
    }

}
